import glob

import cv2
import json
import numpy as np
import torch
from model import mobilenet_v3_s_model


def main(img_path):
    # 读取索引文件
    with open('./index_to_class.json') as f:
        index_to_class = {int(k): v for k, v in json.load(f).items()}

    # 实例化模型
    model = mobilenet_v3_s_model(classes_num=len(index_to_class), mode='predict')
    # model.state_dict()
    # 加载模型权重
    weights = torch.load('./weights/mobilenet_v3_small_mytrain.pth', map_location='cpu')
    model.load_state_dict(weights)

    # 读取一张图片
    img = cv2.imread(img_path)
    h, w = img.shape[0], img.shape[1]
    img = img[int(h * 0.1):int(h * 0.9), int(w * 0.1):int(w * 0.9), :]
    img = cv2.resize(img, (224, 224))  # 缩放
    img = img / 255.0  # 标准化
    img = np.transpose(img, (2, 0, 1))  # 通道转换  # [H, W, C] -> [C, H, W]
    img = np.expand_dims(img, axis=0)  # 添加`batch`维度  # [N, C, H, W]
    img = torch.as_tensor(img, dtype=torch.float32)  # 转换成`tensor`

    # 预测
    model.eval()  # 进入评估模式, 一定要进入评估模式
    op = torch.argmax(model(img), dim=1).item()
    return index_to_class[op], img


if __name__ == '__main__':
    from PIL import Image

    # 1、文件读取
    img_path = glob.glob('./花卉图片/*.png')
    for j, i in enumerate(img_path):
        pre, img = main(i)
        img = Image.open(i)

        # 2、分类保存
        if pre == "daisy":
            img.save(f'./daisy/{str(j)}.jpg')
        elif pre == 'dandelion':
            img.save(f'./dandelion/{str(j)}.jpg')
        elif pre == 'sunflower':
            img.save(f'./sunflower/{str(j)}.jpg')
        else:
            img.save(f'./tulips/{str(j)}.jpg')
